import pandas as pd
import matplotlib.pyplot as plt
import argparse
from collections import defaultdict

def process_papers(input_file, output_dir):
    # Read the CSV file
    df = pd.read_csv(input_file)
    
    # Group papers by year
    papers_by_year = defaultdict(list)
    year_counts = defaultdict(int)
    
    # Process each paper
    for _, row in df.iterrows():
        year = row['year']
        if pd.notna(year) and year != 'unknown':  # Skip unknown years
            year = int(year)
            title = row['title']
            papers_by_year[year].append(title)
            year_counts[year] += 1
    
    # Write papers to files by year
    for year in sorted(papers_by_year.keys()):
        output_file = f"{output_dir}/papers_{year}.txt"
        with open(output_file, 'w', encoding='utf-8') as f:
            f.write(f"Papers published in {year}:\n")
            f.write("=" * 50 + "\n\n")
            for i, title in enumerate(papers_by_year[year], 1):
                f.write(f"{i}. {title}\n")
            f.write(f"\nTotal papers in {year}: {year_counts[year]}")
    
    return year_counts

def plot_papers(year_counts):
    years = sorted(year_counts.keys())
    counts = [year_counts[year] for year in years]
    
    plt.figure(figsize=(10, 6))
    plt.bar(years, counts, color='skyblue')
    plt.xlabel('Year')
    plt.ylabel('Number of Papers')
    plt.title('Number of Papers by Year')
    plt.grid(True, alpha=0.3)
    
    # Set x-axis to show only whole years
    plt.gca().xaxis.set_major_locator(plt.MultipleLocator(1))
    plt.gca().xaxis.set_major_formatter(plt.FormatStrFormatter('%d'))
    
    # Add value labels on top of each bar
    for i, count in enumerate(counts):
        plt.text(years[i], count, str(count), ha='center', va='bottom')
    
    plt.tight_layout()
    plt.savefig('papers_by_year.png')
    plt.close()

def main():
    parser = argparse.ArgumentParser(description='Process and analyze papers data')
    parser.add_argument('--plot', action='store_true', help='Generate a plot of papers by year')
    parser.add_argument('--file', type=str, default='data/orise_papers_by_year.csv',
                      help='Path to the CSV file containing papers data')
    args = parser.parse_args()
    
    # Create output directory if it doesn't exist
    import os
    if not os.path.exists('output'):
        os.makedirs('output')
    
    # Process the papers
    year_counts = process_papers(args.file, 'output')
    
    # Print summary to console
    print("\nSummary of papers by year:")
    print("=" * 30)
    for year in sorted(year_counts.keys()):
        print(f"{year}: {year_counts[year]} papers")
    
    # Generate plot if requested
    if args.plot:
        plot_papers(year_counts)
        print("\nPlot has been saved as 'papers_by_year.png'")

if __name__ == "__main__":
    main()